# -*- coding: utf-8 -*-


from ClassifySocket.classify import ClassifySocketType
import tools

import cv2


def loadFrame(tools, dir):
    src = dir + str(i) + ".jpg"
    flag, frame = tools.readImage(src)
    if not flag:
        exit(-1)
    return frame


def loadCamera(tools):
    flag, frame = tools.loadCamera()
    if not flag:
        exit(-1)
    return frame


# the fellowing functions are used for demos
def getPatches(classify_socket_type, frame, dir="./patches/"):
    try:
        classify_socket_type.getPatches(frame, dir)
    except Exception as e:
        print(e)
        exit(-1)


def setup(classify_socket_type, parameter):
    try:
        classify_socket_type.setup(parameter)
    except Exception as e:
        print(e)
        exit(-1)


def train(classify_socket_type, dir="./patches/", file="parameter", method="OTSU", clusters=2):
    try:
        classify_socket_type.train(
            dir, file, method, clusters)
    except Exception as e:
        print(e)
        exit(-1)


if __name__ == "__main__":
    classify_socket_type = ClassifySocketType()

    setup(classify_socket_type, "./parameter")
    # exit(0)
    classify_socket_type.setClassifyLabel(['A', 'B'])

    train(classify_socket_type, method="KMeans")
    for i in range(2, 14):
        frame = loadFrame(tools, "./data/my_photo-")
        # frame = loadCamera(tools)

        classify_socket_type.run(frame)
        classify_socket_type.draw(frame)
        _, classify = classify_socket_type.getClassify()
        classify_socket_type.clearBuffer()
        print(classify)
        cv2.waitKey(0)

        getPatches(classify_socket_type, frame)
